/**
 * 树
 *
 * @param <T> 树的节点类型
 */
export interface Tree<T> {
  /**
   * 树节点的原始内容
   */
  readonly data: T;

  /**
   * 树的父节点，根节点为空
   */
  readonly parent?: Tree<T>;

  /**
   * 树的深度，从1开始
   */
  readonly level: number;
  /**
   * 树的子节点
   */
  readonly children: Tree<T>[];

  /**
   * 将树转成其他格式的树
   *
   * @param mapper map 函数, 参数 node 传入当前节点，参数 children 传入 map 后的子节点。返回 map 后的数组
   * @return map 函数的返回类型对象
   */
  map<R>(mapper: (node: Tree<T>, children: R[]) => R): R;

  filterMap<R>(mapper: (node: Tree<T>, children: R[]) => R | null): R | null;
}

/**
 * 返回数组的最后一个元素
 */
function last<T>(array: T[]): T {
  return array[array.length - 1];
}

/**
 * 删除数组中最后一个元素，如果非空就抛异常
 * array.pop()!
 */
function removeLast<T>(array: T[]): T {
  return array.pop()!;
}

class TreeImpl<T> implements Tree<T> {
  data: T;

  children: TreeImpl<T>[];

  _parent?: Tree<T>;

  _level?: number;

  public get level(): number {
    return this._level!;
  }

  /**
   * 树的父亲节点，可能为 undefined
   */
  public get parent(): Tree<T> | undefined {
    return this._parent;
  }

  /**
   * 创建一棵树
   *
   * @param node 节点
   * @param children 树的子节点，默认为空
   */
  constructor(node: T, children: TreeImpl<T>[] = []) {
    this.data = node;
    this.children = children;
  }

  map<R>(mapper: (tree: Tree<T>, children: R[]) => R): R {
    return this.filterMap(mapper);
  }

  /**
   * 后续遍历树, 如果 mapper 返回空，则过滤掉这个节点
   */
  filterMap<R>(mapper: (node: Tree<T>, children: R[]) => R | null): R {
    const recordStack = [
      {
        tree: this as Tree<T>,
        index: 0,
      },
    ];
    const resultsStack: R[][] = [[], []];
    while (true) {
      const record = last(recordStack);
      const childrenCount = record.tree.children.length;
      if (record.index !== childrenCount) {
        recordStack.push({
          tree: record.tree.children[record.index],
          index: 0,
        });
        record.index++;
        resultsStack.push([]);
        continue;
      }
      // 当节点的子节点都 map 后，map 当前的节点
      // 此时最后一个 resultsStack 保存着最后一个 recordStack 的树的子节点的被 map 后的结果
      const result = mapper(
        removeLast(recordStack).tree,
        removeLast(resultsStack)
      );
      if (result !== null && result !== undefined) {
        last(resultsStack).push(result);
      }
      if (recordStack.length === 0) {
        break;
      }
    }
    return last(resultsStack).pop()!;
  }
}

export class Trees<T> {
  children: Tree<T>[];

  constructor(children: Tree<T>[]) {
    this.children = children;
  }

  map<U>(mapper: (tree: Tree<T>, children: U[]) => U): U[] {
    return this.children.map((child) => child.map(mapper));
  }

  filterMap<R>(mapper: (node: Tree<T>, children: R[]) => R | null): R[] {
    const results: R[] = [];
    for (let child of this.children) {
      const result = child.filterMap(mapper);
      if (result !== null && result !== undefined) {
        results.push(result);
      }
    }
    return results;
  }
}

/**
 * 将数组转成树(对象)
 *
 * @param list 数组
 * @param getId 获得id的方法
 * @param getParentId 获得父节点的id的方法
 * @return 树的根节点对象
 */
export function parseTree<T, ID>(
  list: T[],
  getId: (element: T) => ID,
  getParentId: (element: T) => ID
): Trees<T> {
  const trees = listToTrees(list, getId, getParentId);
  return new Trees(trees);
}

/**
 * 将数组转成树
 *
 * @param list 数组
 * @param getId 获得id的方法
 * @param getParentId 获得父节点的id的方法
 * @return 树的数组,数组中是所有根节点
 */
export function listToTrees<T, ID>(
  list: T[],
  getId: (element: T) => ID,
  getParentId: (element: T) => ID
): Tree<T>[] {
  if (list.length === 0) {
    return [];
  }
  const map = new Map<ID, TreeImpl<T>>();
  const repeatIdChecker = createRepeatIdChecker();
  for (let element of list) {
    const id = getId(element);
    if (getId(element) === getParentId(element)) {
      throw new TreeError("id equals parentId");
    }
    repeatIdChecker(id);
    map.set(id, new TreeImpl(element));
  }
  const rootNodes: TreeImpl<T>[] = [];
  for (let element of list) {
    const parentId = getParentId(element);
    const id = getId(element);
    const parentTree = map.get(parentId);
    const tree = map.get(id)!;
    if (parentTree != null) {
      parentTree.children.push(tree);
      tree._parent = parentTree;
    } else {
      rootNodes.push(tree);
    }
  }
  if (rootNodes.length === 0) {
    throw new TreeError("none root tree");
  }
  setLevel(rootNodes);
  return rootNodes;
}

class TreeError extends Error {
  name: string;

  constructor(message: string) {
    super(message);
    this.name = "treeError";
  }
}

function createLoopChecker<T>() {
  const set = new Set();
  return (tree: Tree<T>) => {
    if (set.has(tree)) {
      throw new TreeError("loop");
    }
  };
}

function createRepeatIdChecker() {
  const set = new Set();
  return (id) => {
    if (set.has(id)) {
      throw new TreeError(`repeat id ${id}`);
    }
  };
}

function setLevel<T>(nodes: TreeImpl<T>[], level = 1) {
  const loopChecker = createLoopChecker();
  while (nodes.length > 0) {
    const currentLevel = level;
    nodes = nodes.flatMap((tree) => {
      loopChecker(tree);
      tree._level = currentLevel;
      return tree.children;
    });
    level += 1;
  }
}
